#include <ied_vision/ied_vision.h>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <ied_vision/Localization.h>

double IEDVision::colorBound[] = {128,32,32};
double IEDVision::rangeBound = 0.09;
double IEDVision::detectionThreshold = 0.5;
double IEDVision::cameraFieldOfView = 0.523598776;

const int COLOR_UPDATE = 0;
const int DEPTH_UPDATE = 1;

IEDVision::IEDVision()
	: transport(handle)
{
	color_sub = transport.subscribe("color_in", 1, &IEDVision::imageCb, this);
	depth_sub = transport.subscribe("depth_in", 1, &IEDVision::depthCb, this);
	pose_sub = handle.subscribe("pose_in", 1, &IEDVision::poseCb, this);
	ied_pub = handle.advertise< geometry_msgs::Point >("ied_detections");
	
	handle.getParam("rangeBound", rangeBound);
	handle.getParam("blueBound", colorBound[0]);
	handle.getParam("greenBound", colorBound[1]);
	handle.getParam("redBound", colorBound[2]);
	handle.getParam("detectionThreshold", detectionThreshold);
	handle.getParams("cameraFOV", cameraFieldOfView);
	
	gridParams = Localization_Env_Read(); //Replace with getParams later
	iedLocations = Localization(gridParams.GetXGridSize(), gridParams.GetYGridSize(),
	                            gridParams.GetXSquareSize(), gridParams.GetYSquareSize(),
	                            gridParams.GetDeadZoneSize(), gridParams.GetPathWidthSize(),
	                            gridParams.GetAngleOffset());
	
	updated[0] = false;
	updated[1] = false;
	
#ifdef DEBUG1
	cv::namedWindow(WINDOW1);
#endif
#ifdef DEBUG2
	cv::namedWindow(WINDOW2);
	cv::namedWindow(WINDOW2A);
#endif
}

IEDVision::~IEDVision()
{
#ifdef DEBUG1
    cv::destroyWindow(WINDOW1);
#endif
#ifdef DEBUG2
    cv::destroyWindow(WINDOW2);
	cv::destroyWindow(WINDOW2A);
#endif
}

void IEDVision::imageCb(const sensor_msgs::ImageConstPtr& msg)
{
	cv_bridge::CvImageConstPtr cv_ptr;
	try
	{
		cv_ptr = cv_bridge::toCvShare(msg, enc::BGR8);
	}
	catch (cv_bridge::Exception& e)
	{
		ROS_ERROR("cv_bridge exception: %s", e.what());
		return;
	}
	
	cv::Mat colorChannels[3];
	cv::split(cv_ptr->image, colorChannels);
	for(int i=0; i<COLOR_CHANNELS; ++i)
	{
		cv::Mat rawEdges;
		cv::Sobel(colorChannels[i], rawEdges, CV_16SC1, 1, 0);
		cv::compare(rawEdges, colorBound[i], colorEdges[i], cv::CMP_GE);
	}
	
	updated[COLOR_UPDATE] = true;
	findIEDs();

#ifdef DEBUG1
	cv::imshow(WINDOW1, colorEdges[0]);
	cv::waitKey(3);
#endif
}

void IEDVision::depthCb(const sensor_msgs::ImageConstPtr& msg)
{
	cv_bridge::CvImageConstPtr cv_ptr;
	try
	{
		cv_ptr = cv_bridge::toCvShare(msg);
	}
	catch (cv_bridge::Exception& e)
	{
		ROS_ERROR("cv_bridge exception: %s", e.what());
		return;
	}

	rangeImage = cv_ptr->image;

	cv::Mat rawEdges;
	cv::Sobel(cv_ptr->image, rawEdges, CV_32FC1, 1, 1);
	cv::compare(rawEdges, rangeBound, rangeEdges, cv::CMP_GE);
	
	updated[DEPTH_UPDATE] = true;
	findIEDs();

#ifdef DEBUG2
	cv::imshow(WINDOW2, rangeEdges);
	cv::imshow(WINDOW2A, rawEdges);
	cv::waitKey(3);
#endif
}

void IEDVision::findIEDs()
{
	//Only process if we have all new (hopefully synchronized) data
	if(!updated[DEPTH_UPDATE] || !updated[COLOR_UPDATE])
	{
		return;
	}
	
	updated[DEPTH_UPDATE] = false;
	updated[COLOR_UPDATE] = false;
	
	//This vector may be replaced by a call to the localization class to get each IED's location
	for(int x = 0; x < gridParams.GetXGridSize(); ++x)
	{
		for(int y = 0; y < gridParams.GetYGridSize(); ++y)
		{
			SquareInfo iedLocale = iedLocations.GetSquareInfo(x, y);
			
			// Adjust the angle to be in a coordinate system with 0 on the left side of the image,
			// with angle increasing to the right
			double lowZeroed = 0.0;
			double highZeroed = 0.0;
			if(iedLocale._lowAngle >= M_PI)
				lowZeroed = (2*M_PI - iedLocale._lowAngle) + (cameraFieldOfView/2);
			else
				lowZeroed = (cameraFieldOfView/2) - iedLocale._lowAngle;
			if(iedLocale._highAngle >= M_PI)
				highZeroed = (2*M_PI - iedLocale._highAngle) + (cameraFieldOfView/2);
			else
				highZeroed = (cameraFieldOfView/2) - iedLocale._highAngle;
			
			if((lowZeroed >= 0 && lowZeroed <= cameraFieldOfView) ||
			   (highZeroed <= cameraFieldOfView && highZeroed >= 0)
			{
				cv::Mat tempMask;
				cv::Mat iedMask;
				
				// Highlight those portions of the image in the possible IED squqre ranges.
				// Dilate a bit for edge effects.
				cv::inRange(rangeImage, cv::Scalar(iedLocale._lowRange), cv::Scalar(iedLocale._highRange), tempMask);
				cv::dilate(tempMask, iedMask, cv::Mat());
				
				// Correlate angles to pixels, and if out of view pin to extreme edge.
				int vertexPixelOne = rangeImage.cols*(lowZeroed/cameraFieldOfView);
				vertexPixelOne = vertexPixelOne < 0 ? 0 : vertexPixelOne;
				vertexPixelOne = vertexPixelOne > rangeImage.cols ? rangeImage.cols : vertexPixelOne;
				int vertexPixelTwo = rangeImage.cols*(highZeroed/cameraFieldOfView);
				vertexPixelOne = vertexPixelTwo < 0 ? 0 : vertexPixelTwo;
				vertexPixelOne = vertexPixelTwo > rangeImage.cols ? rangeImage.cols : vertexPixelTwo;
				// Create a mask only active in the IEDs possible anglular range
				cv::Mat angleMask(rangeImage.rows, rangeImage.cols, cv::CV_8U, cv::Scalar(0));
				cv::rectangle(angleMask, cv::Point(vertexPixelOne,0), 
				              cv::Point(vertexPixelTwo,rangeImage.rows), 
				              cv::Scalar(-1), cv::CV_FILLED);
				
				// Figure the max number of points that could be active due to range edges
				cv::bitwise_and(iedMask, angleMask, tempMask);
				int iedMaskPoints = cv::countNonZero(tempMask);
			
				// Find number of color/range edges that lie in range...
				cv::Mat foundEdges;
				int matchingPoints[COLOR_CHANNELS+1];
				int maxPoints[COLOR_CHANNELS+1];
				for(int i=0; i<COLOR_CHANNELS; ++i) 
				{
					cv::bitwise_and(colorEdges[i], angleMask, tempMask); //Remove portions not at IED angles
					cv::bitwise_and(iedMask, tempMask, foundEdges); //Remove portions not in range
					matchingPoints[i] = cv::countNonZero(foundEdges);
					maxPoints[i] = std::min(cv::countNonZero(tempMask), iedMaskPoints);
				}
				cv::bitwise_and(rangeEdges, angleMask, tempMask);
				cv::bitwise_and(iedMask, tempMask, foundEdges);
				matchingPoints[COLOR_CHANNELS] = cv::countNonZero(foundEdges);
				maxPoints[COLOR_CHANNELS] = std::min(cv::countNonZero(tempMask), iedMaskPoints);
		
				// Score by summing the ratio of active pixels to total possible that
				// could be active for each plane. 
				float totalScore = 0.0f;
				for(int i=0; i<COLOR_CHANNELS+1; ++i)
				{
					totalScore += (float)matchingPoints[i]/maxPoints[i];
				}
				if(totalScore/(COLOR_CHANNELS+1) >= detectionThreshold/(COLOR_CHANNELS+1))
				{
					geometry_msgs::Point outPnt;
					outPnt.x = x;
					outPnt.y = y;
					outPnt.z = 0;
					ied_pub.publish(outPnt);
				}
			}
		}
	}
	
}

void IEDVision::poseCb(const geometry_msgs::PoseStampedConstPtr& msg)
{
	tf::Stamped< tf::Pose > pose;
	tf::poseStampedMsgToTF(*msg, pose);
	
	iedLocations.Update(pose.getOrigin().x(), pose.getOrigin().y(), 
	                    pose.getOrigin().z(), pose.getRotation().getAngle());
}

int main(int argc, char** argv)
{
	ros::init(argc, argv, "ied_vision");
	IEDVision iv;
	ros::spin();
	return 0;
}
